using Random, Statistics, Printf, Plots

# ===============================
# 邻居表函数保持不变
function neighbor(n)
    nbor = zeros(Int, n * n, 8)
    for ispin in 1:n*n
        iy = div(ispin - 1, n) + 1
        ix = mod(ispin - 1, n) + 1
        ixp = ix < n ? ix + 1 : 1
        iyp = iy < n ? iy + 1 : 1
        ixm = ix > 1 ? ix - 1 : n
        iym = iy > 1 ? iy - 1 : n
        nbor[ispin, 1] = (iy - 1) * n + ixp
        nbor[ispin, 2] = (iyp - 1) * n + ix
        nbor[ispin, 3] = (iy - 1) * n + ixm
        nbor[ispin, 4] = (iym - 1) * n + ix
        nbor[ispin, 5] = (iyp - 1) * n + ixp
        nbor[ispin, 6] = (iyp - 1) * n + ixm
        nbor[ispin, 7] = (iym - 1) * n + ixm
        nbor[ispin, 8] = (iym - 1) * n + ixp
    end
    return nbor
end

function neighbor_SS(n)
    nbor = neighbor(n)
    nbor1 = zeros(Int, n * n * 4, 5)
    for i in 1:n*n*4
        A = div((i - 1), 4) + 1
        a = i - 4 * (A - 1)
        if a == 1
            nbor1[i, 1] = i+1
            nbor1[i, 2] = i+2
            nbor1[i, 3] = (nbor[A,3]-1)*4+2
            nbor1[i, 4] = (nbor[A,4]-1)*4+3
            nbor1[i, 5] = (nbor[A,3]-1)*4+4
        elseif a == 2
            nbor1[i, 1] = (nbor[A,1]-1)*4+1
            nbor1[i, 2] = i+2
            nbor1[i, 3] = i-1
            nbor1[i, 4] = (nbor[A,4]-1)*4+4
            nbor1[i, 5] = (nbor[A,4]-1)*4+3
        elseif a == 3
            nbor1[i, 1] = i+1
            nbor1[i, 2] = (nbor[A,2]-1)*4+1
            nbor1[i, 3] = (nbor[A,3]-1)*4+4
            nbor1[i, 4] = i-2
            nbor1[i, 5] = (nbor[A,2]-1)*4+2
        else
            nbor1[i, 1] = (nbor[A,1]-1)*4+3
            nbor1[i, 2] = (nbor[A,2]-1)*4+2
            nbor1[i, 3] = i-1
            nbor1[i, 4] = i-2
            nbor1[i, 5] = (nbor[A,1]-1)*4+1
        end
    end
    return nbor1
end

initspin_bc(n::Int) = rand([-1, 1], n*n*4)

# ===============================
# 改进的能量计算（添加验证）
function energy_total(spin, nbor1, J1, J2, h=1)
    E_exchange = 0.0
    M = 0.0
    N = length(spin)
    
    for i in 1:N
        # 验证邻居索引有效性
        for k in 1:5
            nb = nbor1[i,k]
            if nb < 1 || nb > N
                error("Invalid neighbor index: i=$i, k=$k, nb=$nb")
            end
        end
        
        # 交换作用能
        for k in 1:4
            E_exchange += J1 * spin[i] * spin[nbor1[i,k]]
        end
        E_exchange += J2 * spin[i] * spin[nbor1[i,5]]
        M += spin[i]
    end
    
    E_total = E_exchange / 2.0 - h * M
    
    # 验证能量范围合理性
    if abs(E_total) > 1000
        @warn "Suspicious energy value: E = $E_total"
    end
    
    return E_total
end

function deltaE(ns, spin, nbor1, J1, J2, h=1)
    ssum = 0.0
    for k in 1:4
        ssum += J1 * spin[nbor1[ns,k]]
    end
    ssum += J2 * spin[nbor1[ns,5]]
    return -2.0 * spin[ns] * ssum + 2.0 * h * spin[ns]
end

# ===============================
# 改进的能级离散化
function generate_levels_2d(n, J1, J2, h=1)
    Nspin = n*n*4
    
    # 更精细的能量步长
    min_coupling = min(abs(J1), abs(J2), abs(h))
    stepE = max(0.5, min_coupling)  # 避免步长太小
    
    # 扩展能量范围，考虑所有可能状态
    Emin_possible = - (4*abs(J1) + abs(J2)) * Nspin/2 - abs(h) * Nspin
    Emax_possible = -Emin_possible
    
    # 确保范围足够大
    range_extension = 2.0 * max(abs(J1), abs(J2), abs(h)) * Nspin
    Emin = min(Emin_possible, -range_extension)
    Emax = max(Emax_possible, range_extension)
    
    Energy = collect(Emin:stepE:Emax)
    Magnet = collect(-Nspin:2:Nspin) .|> Float64
    
    println("Energy levels: $(length(Energy)), Magnet levels: $(length(Magnet))")
    println("Energy range: $Emin to $Emax, step: $stepE")
    
    return Energy, Magnet
end

@inline function em_index(E, M, Energy, Magnet)
    e_idx = round(Int, (E - Energy[1]) / (Energy[2] - Energy[1])) + 1
    m_idx = round(Int, (M - Magnet[1]) / 2.0) + 1
    e_idx = clamp(e_idx, 1, length(Energy))
    m_idx = clamp(m_idx, 1, length(Magnet))
    return e_idx, m_idx
end

# 改进的平坦性判断
function isflat_2d(Hist; flat_crit=0.08, flat_frac=0.8)
    # 只考虑被访问过的bin
    visited = findall(Hist .> 0)
    if length(visited) < 10  # 至少需要10个被访问的bin
        return false
    end
    
    vals = Hist[visited]
    avg = mean(vals)
    if avg == 0
        return false
    end
    
    # 相对标准差作为平坦性判断
    relative_std = std(vals) / avg
    return relative_std < flat_crit
end

# ===============================
# 改进的WL采样
function wang_landau_SS_2d(n, J1, J2, h=1;
        mstep=20_000_000, min_lnf=1e-12,  # 更多步数，更严格收敛
        flat_crit=0.08, flat_frac=0.8, check_factor=50)

    Nspin = n*n*4
    nbor1 = neighbor_SS(n)
    spin = initspin_bc(n)

    Energy, Magnet = generate_levels_2d(n, J1, J2, h)
    NE, NM = length(Energy), length(Magnet)

    S = zeros(Float64, NE, NM)
    Hist = zeros(Int, NE, NM)

    # 初始状态验证
    Ecurr = energy_total(spin, nbor1, J1, J2, h)
    Mcurr = sum(spin)
    e_idx, m_idx = em_index(Ecurr, Mcurr, Energy, Magnet)
    
    println("Initial: E=$Ecurr, M=$Mcurr, index=($e_idx, $m_idx)")

    ln_f = 1.0
    check_every = max(100, check_factor * Nspin)  # 更频繁检查
    iter = 1
    flat_count = 0
    max_flat_count = 3  # 连续平坦3次才认为收敛

    println("=== WL stage (E,M) with h=$h ===")
    println("System size: $Nspin spins")
    println("State space: $NE × $NM = $(NE*NM) bins")

    for step in 1:mstep
        ns = rand(1:Nspin)
        de = deltaE(ns, spin, nbor1, J1, J2, h)
        dm = -2.0 * spin[ns]
        E_new = Ecurr + de
        M_new = Mcurr + dm
        e_new, m_new = em_index(E_new, M_new, Energy, Magnet)

        # 边界检查
        if e_new < 1 || e_new > NE || m_new < 1 || m_new > NM
            @warn "State out of bounds: E_new=$E_new, M_new=$M_new"
            continue
        end

        ΔS = S[e_idx, m_idx] - S[e_new, m_new]
        if rand() < (ΔS >= 0 ? 1.0 : exp(ΔS))
            spin[ns] = -spin[ns]
            Ecurr, Mcurr = E_new, M_new
            e_idx, m_idx = e_new, m_new
        end
        
        S[e_idx, m_idx] += ln_f
        Hist[e_idx, m_idx] += 1

        # 进度报告
        if step % (10*check_every) == 0
            visited_bins = count(Hist .> 0)
            coverage = visited_bins / (NE * NM)
            @printf("Step %d: ln_f=%.2e, coverage=%.1f%%, visited=%d/%d\n",
                   step, ln_f, 100*coverage, visited_bins, NE*NM)
        end

        if step % check_every == 0
            if isflat_2d(Hist; flat_crit=flat_crit, flat_frac=flat_frac)
                flat_count += 1
                @printf("Iter %d flat (%d/%d) @ step %d ln_f=%.6e\n",
                       iter, flat_count, max_flat_count, step, ln_f)
                
                if flat_count >= max_flat_count
                    ln_f *= 0.5
                    fill!(Hist, 0)
                    iter += 1
                    flat_count = 0
                    
                    if ln_f < min_lnf
                        println("WL converged after $iter iterations.")
                        break
                    end
                end
            else
                flat_count = 0  # 重置连续平坦计数
            end
        end
    end

    # 归一化态密度
    S .-= minimum(S)
    
    # 检查态密度覆盖范围
    visited = count(S .> 0)
    println("Final: visited $visited bins out of $(NE*NM)")
    
    return S, Energy, Magnet
end

# 后处理函数保持不变
function thermo_from_S_EM(S, Energy, Magnet; Nspin::Int)
    Ts = range(0.1, 8, length=100)
    Eavg=zeros(length(Ts)); Cv=zeros(length(Ts))
    Mavg=zeros(length(Ts)); Mabs=zeros(length(Ts))
    
    for (tidx,T) in enumerate(Ts)
        β = 1/T
        maxlogw = -Inf
        for ei in eachindex(Energy), mi in eachindex(Magnet)
            lw = S[ei, mi] - β * Energy[ei]
            if lw > maxlogw; maxlogw = lw; end
        end
        
        Z = 0.0; E1 = 0.0; E2 = 0.0
        M1 = 0.0; Mabs1 = 0.0
        
        for ei in eachindex(Energy), mi in eachindex(Magnet)
            if S[ei, mi] > 0  # 只考虑被访问的状态
                w = exp(S[ei, mi] - β * Energy[ei] - maxlogw)
                Z += w
                E1 += Energy[ei] * w
                E2 += Energy[ei]^2 * w
                M1 += Magnet[mi] * w
                Mabs1 += abs(Magnet[mi]) * w
            end
        end
        
        if Z > 0
            Eavg[tidx] = (E1/Z) / Nspin
            Cv[tidx] = (β^2 * (E2/Z - (E1/Z)^2)) / Nspin
            Mavg[tidx] = (M1/Z) / Nspin
            Mabs[tidx] = (Mabs1/Z) / Nspin
        end
    end
    
    return Ts, Eavg, Cv, Mavg, Mabs
end

# ===============================
# 主程序：添加验证和比较
function main()
    Random.seed!(1234)
    
    # 尝试不同的系统尺寸
    for n in [2]
        println("\n" * "="^50)
        println("Testing n = $n (Nspin = $(n*n*4))")
        println("="^50)
        
        J1, J2 = 1.0, 2.0
        h_target = 1.0

        @time S2d, Elist, Mlist = wang_landau_SS_2d(n, J1, J2, h_target)

        Ts, Eavg, Cv, Mavg, Mabs = thermo_from_S_EM(S2d, Elist, Mlist; Nspin=n*n*4)

        # 输出关键物理量
        println("\nKey results for n=$n, h=$h_target:")
        println("Peak C_v at T = $(Ts[argmax(Cv)])")
        println("Magnetization at T=0.1: $(Mabs[1])")
        # 打印结果表
       println("===== h = $h 的热力学量 =====")
       @printf("%8s %15s %15s %15s %15s\n", "T", "E/N", "C_v", "M/N", "|M|/N")
       for i in eachindex(Ts)
       @printf("%8.4f %15.6f %15.6f %15.6f %15.6f\n",
        Ts[i], Eavg[i], Cv[i], Mavg[i], Mabs[i])
       end

        # 绘图
        p1 = plot(Ts, Eavg, xlabel="T", ylabel="E/N", title="Energy (n=$n)", lw=2)
        p2 = plot(Ts, Cv, xlabel="T", ylabel="C_v", title="Specific Heat (n=$n)", lw=2)
        p3 = plot(Ts, Mabs, xlabel="T", ylabel="|M|/N", title="Magnetization (n=$n)", lw=2)
        display(plot(p1, p2, p3, layout=(3,1), size=(800,600)))
    end
end

# 运行主程序
main()